import os
from dataclasses import dataclass
from functools import lru_cache, partial
from typing import Any, Union, Callable, Tuple, TypeVar, Dict, List

import mne

from util import *
from src.hook import hookdata_decorator
mne.set_log_level('WARNING')

Event = TypeVar('Event', bound=Tuple[np.ndarray, int, int, int, str])

@dataclass
class WindowInfo:
    start:int
    end:int
    label:str
    nId:int = 0
    spec:str = ''

def read_edf(path, new_freq:int=100, l_freq:int=1, h_freq:int=30, 
             verbose=True) -> Tuple[mne.io.Raw, float]:
    '''
    功能：读取edf文件
    参数：
        path：edf文件路径
        new_freq：降采样采样率
    返回：预处理后的raw对象，SRR（旧采样率/新采样率）
    '''  
    raw:mne.io.Raw = mne.io.read_raw(path, preload=True)
    SRR = raw.info['sfreq']/new_freq
    raw = raw.resample(new_freq, npad='auto').filter(l_freq=l_freq, h_freq=h_freq)
    if verbose:
        print(f'raw: resample({new_freq})|filter({l_freq}, {h_freq})')
        # print(f'  \u21AA SRR: {SRR}')
        print(f'  -> SRR: {SRR}')
    return raw, SRR

def read_tags(tags_path, SRR:Union[int, float]=1, **kwargs) -> Iterator[WindowInfo]:
    '''
    功能：读取tags文件事件信息
    参数：
        tags_path：tags文件路径
    '''
    event_data = loadJSON(tags_path)['EventData']
    for event in event_data:
        start, duration, event_name = event['Start'], event['Duration'], event['EventName']
        start, duration = map(lambda x:int(x/SRR), (start, duration))
        yield WindowInfo(start, duration, event_name)  # 不分割

def window_offset(start:int, end:int, direct:str, method:str, dist:int, step:int, n_times:int):
    '''
    功能：窗口偏移
    参数：
        start：开始时间
        end：结束时间
        dist：差距
        direct：偏移方向
        method：偏移方法
    '''
    turn = 0 if method == 'fill' else 1
    if dist > 0:
        offset = (1-turn)*(step-dist%step)-turn*dist%step  # 滑动间隔
    else: 
        offset = (1-turn)*(-dist)-turn*(end-start)  # 窗口长度
    if direct == 'f':
        start -= offset
    elif direct == 'b':
        end += offset
    else:
        start -= offset//2
        end += offset//2 + offset%2
    if method == 'fill':
        offset = step if dist > 0 else end-start
        start = start+offset if start < 0 else start
        end = end-offset if end > n_times else end
    return start, end

@id_decorator(kwarg_keys=('mode',))
def window_split(src:Iterator[WindowInfo], n_times:int=float('inf'), window:int=100, step:int=50, 
                 mode:str='mlack') -> Iterator[WindowInfo]:
    '''
    功能：划分窗口
    参数：
        n_times：时间长度
        window：窗口长度
        step：滑动间隔
        mode：f/b/m-前缀，表示前中后，fill/lack-词根，填充/缺失
        lack：缺窗口，flack前缺，black后缺，mlack前后缺
    返回：（窗口开始时间，窗口结束时间，窗口类型）
    '''
    direct, method = mode[0], mode[1:]
    check_param((direct, 'fbm'), (method, ['fill', 'lack']))
    
    class WindowForm:
        normal = 0
        multi = 1
        single = 2

    for wi in src:
        start, end = wi.start, wi.end
        dist = end - window - start
        form = WindowForm.normal  # 无缺
        if dist != 0:
            form = WindowForm.multi if dist > 0 else WindowForm.single  # 多窗口缺，单窗口缺
            start, end = window_offset(start, end, direct, method, dist, step, n_times)
        for i in range(start, end-window+1, step):
            yield WindowInfo(i, i + window, wi.label, spec=f'form{form}')
        if method == 'lack' and form == WindowForm.single:
            yield wi

def jaccard_similarity(pos1:Tuple[int, int], pos2:Tuple[int, int]) -> float:
    '''杰卡德相似系数'''
    set1 = set(range(pos1[0], pos1[1]))
    set2 = set(range(pos2[0], pos2[1]))
    return len(set1.intersection(set2))/len(set1.union(set2))

def get_search_area(windows:List[WindowInfo], sp:int=0, window_len:int=100
                    ) -> Tuple[int, int, List[WindowInfo], int]:
    '''获取搜索区域'''
    if sp == len(windows):
        return float('inf'), float('inf'), [], sp
    start, end, search_obj = windows[sp].start, 0, []
    for i in range(sp, len(windows)):
        wi = windows[i]
        search_obj.append(wi)
        end = max(end, wi.end)
        if end - start >= window_len or i == len(windows)-1:
            return start, end, search_obj, i + 1

@id_decorator(kwarg_keys=('mode',))
def window_split_standard(src:Iterator[WindowInfo], n_times:int=float('inf'), window:int=100, step:int=50, 
                          mode:str='max', threshold:int=0) -> Iterator[WindowInfo]:
    '''
    功能：分割窗口的标准方法
    参数：
        mode：max/last/first，最大/最后/最前，标签选取方式
        threshold：分数阈值
    '''
    src:list = list(src)
    src.sort(key=lambda x:x.start)  # 按起始时间排序
    search_area = get_search_area(src, sp=0, window_len=window)
    area_start, area_end, search_obj, sp = search_area
    for i in range(0, n_times-window+1, step):
        start, end = i, i + window
        out_event_name = '正常'
        spec = f'unmarked'
        scores = []
        if end > area_start:  # 接近搜索区域
            for wi in search_obj:
                sorce = jaccard_similarity((wi.start, wi.end), (start, end))
                scores.append(sorce)
            temp = [s for s in scores if s > threshold]
            if temp:
                if mode == 'last':
                    score = temp[-1]
                elif mode == 'first':
                    score = temp[0]
                else:
                    score = max(temp)
                wi = search_obj[scores.index(score)]
                spec = join_vars('marked', src.index(wi), to_percent(score), sep='&')
                out_event_name = wi.label
        if start > area_end:  # 远离搜索区域, 更新下一个搜索区域
            search_area = get_search_area(src, sp=sp, window_len=window)
            area_start, area_end, search_obj, sp = search_area
        yield WindowInfo(start, end, out_event_name, spec=spec)

def parse_pattern(pattern:str) -> Tuple[list, Any, bool]:
    '''
    功能：解析规则
    参数：
        pattern：规则
    返回：（映射前值，映射后值，选择器）
    '''
    pattern = pattern.replace(' ', '')
    fval, bval = pattern.split('->')

    _not = False
    if fval.startswith('^'):
        _not = True
        fval = fval[1:]
    if fval.startswith('[') and fval.endswith(']'):
        fval = fval[1:-1]
        if fval.startswith('^'):
            _not = True
            fval = fval[1:]
        fval = list(fval)
    if isinstance(fval, str):
        temp = fval.split(',')
        fval = temp if temp else [fval]

    fval = list(map(lambda x:eval(x), fval))
    bval = eval(bval)
    logic = lambda x, y: x not in y if _not else x in y
    selector = lambda x, y, z: z if logic(x, y) else x
    selector = partial(selector, y=fval, z=bval)
    return selector

def val_mapping(obj:Any, *patterns:str) -> Any:
    '''
    功能：值映射
    参数：
        obj：映射前对象
        *patterns：规则
    返回：映射后对象
    '''
    for pattern in patterns:
        selector = parse_pattern(pattern)
        if isinstance(obj, dict):
            mapping = lambda x: {k:selector(v) for k, v in x.items()}
            obj = mapping(obj)
        else:
            raise ValueError(f'Unsupported type: {type(obj)}')
    return obj

@lru_cache
def get_label_map() -> Dict[str, str]:
    '''
    功能：获取标签映射字典
    '''
    config = loadYAML('config.yaml')
    label_map:dict = config['label_map'][1]
    label_map = val_mapping(label_map, '1 -> 2', '11 -> 1', '2 -> 0', '[^01] -> 2')
    return label_map

@hookdata_decorator()
def padding_data(data:np.ndarray, target_len:int, axis:int=0) -> np.ndarray:
    '''
    功能：填补数据
    参数：
        data：数据
        target_len：目标长度
        threshold：填补阈值
        axis：轴方向
    '''
    tdata = np.moveaxis(data, axis, -1)
    ori_len = tdata.shape[0]
    if ori_len >= target_len:
        return data
    x_old = np.arange(ori_len)
    x_new = np.linspace(0, ori_len, target_len)
    new_data = np.empty((target_len, tdata.shape[-1]), dtype=data.dtype)
    for i in range(tdata.shape[-1]):
        new_data[:, i] = np.interp(x_new, x_old, tdata[:, i])
    return np.moveaxis(new_data, -1, axis)

def _event_generator(raw:mne.io.Raw, window_prod:Callable[..., Iterator[WindowInfo]], 
                     window_len:int=100, pad_mode='interp', verbose=True) -> Iterator[Event]:
    '''
    功能：生成事件数据
    参数：
        raw：raw对象
        window_prod：窗口生成者
    返回：（数据，开始时间，结束时间，编码标签，窗口id）
    '''
    ori_data = raw.get_data()[:16]  # 16个电极
    n_times = raw.n_times
    label_map = get_label_map()
    try:
        for wi in window_prod(n_times=n_times):
            start, end = wi.start, wi.end
            label = label_map[wi.label]
            if end - start < window_len:
                if pad_mode == 'interp':
                    data = padding_data(ori_data[:, start:end], window_len, axis=0)
                    wi.spec = join_vars(wi.spec, f'pad{end-start}', sep='&', non_null=True)
                elif pad_mode == 'ffill':
                    start = end - window_len
                    if start >= 0:
                        data = ori_data[:, start:end]
                    else: continue
                elif pad_mode == 'bfill':
                    end = start + window_len
                    if end <= n_times:
                        data = ori_data[:, start:end]
                    else: continue
                else:
                    data = ori_data[:, start:end]
            else:
                data = ori_data[:, start:end]
            wid = join_vars(wi.start, wi.nId, wi.spec, non_null=True)
            if verbose:
                print(f'window: {start}-{end}, {label}, {wid}')
            yield data, start, end, label, wid
    finally:
        raw.close()

def eeg_preprocess(*paths, mode='test', split_mode=None, verbose=True) -> Iterator[Tuple[str, float, Iterator[Event]]]:
    '''
    功能：eeg数据预处理
    参数：
        *paths：路径
        verbose：是否输出调试信息
    返回：（项目名，采样比，事件生成器）
    '''
    edf_paths = []
    for path in paths:
        if os.path.isdir(path):
            edf_paths += find_edf_paths(path)
        elif os.path.isfile(path):
            edf_paths.append(path)
        else:
            raise ValueError(f'{path} does not exist.')
    for i, edf_path in enumerate(edf_paths):
        tags_path = find_tags_path(edf_path)
        if tags_path is None:
            if verbose:
                print(f'{tags_path} is not exist')
            continue
        raw, SRR = read_edf(edf_path, verbose=verbose)
        
        if mode == 'raw':
            window_prod = id_decorator()(partial(read_tags, tags_path=tags_path, SRR=SRR))
            gen = _event_generator(raw, window_prod, pad_mode=None, verbose=verbose)
        elif mode == 'test':
            src = read_tags(tags_path=tags_path, SRR=SRR)
            window_prod = partial(window_split_standard, src=src, mode=split_mode)
            gen = _event_generator(raw, window_prod, verbose=verbose)
        elif mode == 'train':
            src = read_tags(tags_path=tags_path, SRR=SRR)
            window_prod = partial(window_split, src=src, mode=split_mode)
            gen = _event_generator(raw, window_prod, verbose=verbose)
        else:
            raise ValueError(f'The mode is not supported: {mode}')
        
        yield f'{i}_{filename_base(edf_path)}', SRR, gen
        if verbose:
            print('-'*30)

if __name__ == '__main__':
    fz_folder = '../data/fz'
    for fbname, SRR, generator in eeg_preprocess(fz_folder, mode='train'):
        for t_eeg, start, end, label, wid in generator:
            pass
